/*
 * Copyright 2019 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
@file:OptIn(ExperimentalStdlibApi::class)

package com.gitee.wsl.flow.multicast.channel

import com.gitee.wsl.flow.multicast.SharedFlowProducer
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.channels.SendChannel
import kotlin.coroutines.cancellation.CancellationException

internal interface ChannelManager<T> {

    suspend fun addDownstream(
        channel: SendChannel<Message.Dispatch.Value<T>>,
        piggybackOnly: Boolean = false
    )

    suspend fun removeDownstream(channel: SendChannel<Message.Dispatch.Value<T>>)

    suspend fun close()

    /**
     * Holder for each downstream collector
     */
    data class ChannelEntry<T>(
        /**
         * The channel used by the collector
         */
        private val channel: SendChannel<Message.Dispatch.Value<T>>,
        /**
         * Tracking whether this channel is a piggyback only channel that can be closed without ever
         * receiving a value or error.
         */
        val piggybackOnly: Boolean = false
    ) {
        private var _awaitsDispatch: Boolean = !piggybackOnly

        val awaitsDispatch
            get() = _awaitsDispatch

        suspend fun dispatchValue(value: Message.Dispatch.Value<T>) {
            _awaitsDispatch = false
            try {
                channel.send(value)
            } catch (e: CancellationException) {
                // ignore
            }
        }

        fun dispatchError(error: Throwable) {
            _awaitsDispatch = false
            channel.close(error)
        }

        fun close() {
            channel.close()
        }

        fun hasChannel(channel: SendChannel<Message.Dispatch.Value<T>>) = this.channel === channel

        fun hasChannel(entry: ChannelEntry<T>) = this.channel === entry.channel
    }

}

/**
 * Messages accepted by the [ChannelManager].
 */
sealed class Message<out T> {
    /**
     * Add a new channel, that means a new downstream subscriber
     */
    class AddChannel<T>(
        val channel: SendChannel<Dispatch.Value<T>>,
        val piggybackOnly: Boolean = false
    ) : Message<T>()

    /**
     * Remove a downstream subscriber, that means it completed
     */
    class RemoveChannel<T>(val channel: SendChannel<Dispatch.Value<T>>) : Message<T>()

    sealed class Dispatch<out T> : Message<T>() {
        /**
         * Upstream dispatched a new value, send it to all downstream items
         */
        class Value<out T>(
            /**
             * The value dispatched by the upstream
             */
            val value: T,
            /**
             * Ack that is completed by all receiver. Upstream producer will await this before asking
             * for a new value from upstream
             */
            val delivered: CompletableDeferred<Unit>
        ) : Dispatch<T>()

        /**
         * Upstream dispatched an error, send it to all downstream items
         */
        class Error(
            /**
             * The error sent by the upstream
             */
            val error: Throwable
        ) : Dispatch<Nothing>()

        class UpstreamFinished<T>(
            /**
             * SharedFlowProducer finished emitting
             */
            val producer: SharedFlowProducer<T>
        ) : Dispatch<T>()
    }
}


internal fun <T> Message.Dispatch.Value<T>.markDelivered() =
    delivered.complete(Unit)
